from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.vector_stores.milvus import MilvusVectorStore
import asyncio
import os
from datetime import datetime

from llama_index.core.node_parser import SentenceSplitter
from llama_index.core import SimpleDirectoryReader, Document
from abc import abstractmethod
from llama_index.core import VectorStoreIndex, load_index_from_storage, SummaryIndex
from llama_index.core.indices.base import BaseIndex
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.response_synthesizers import ResponseMode
from llama_index.core.storage.storage_context import DEFAULT_PERSIST_DIR, StorageContext


from llama_index.vector_stores.milvus import MilvusVectorStore
from pymilvus import connections, Collection, CollectionSchema, FieldSchema, DataType

from llama_index.core import Settings
from embeddings import embed_model_local_bge_small
from llms import deepseek_llm
Settings.embed_model = embed_model_local_bge_small()
Settings.llm = deepseek_llm()

from config import RagConfig


vector_store = MilvusVectorStore(
    uri=RagConfig.milvus_uri,
    collection_name='bbbb', dim=RagConfig.embedding_model_dim, overwrite=False
)
index=VectorStoreIndex.from_vector_store(vector_store=vector_store)

retriever = index.as_retriever(similarity_top_k=2)
nodes = retriever.retrieve("长春的体育场有哪些？")
print(len(nodes))
print(nodes)
